from dataset import *

import torch.nn as nn


# 编码器 Encoder
class Encoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, hidden_size, num_layers=1, dropout=0.1):
        super(Encoder, self).__init__()
        # nn.Embedding：PyTorch的嵌入层模块
        # vocab_size：词汇表大小，决定嵌入矩阵的行数
        # embedding_size：嵌入向量维度，决定嵌入矩阵的列数
        # padding_idx：指定填充符号PAD的索引，该位置的嵌入向量不会被更新
        self.embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=src_vocab.stoi[PAD])
        # 接收词嵌入维度作为输入大小
        # 输出指定的隐藏层大小
        # 设置LSTM层数（num_layers）
        # 采用batch_first格式处理数据
        # 根据层数决定是否使用dropout：单层时不使用，多层时使用指定的dropout率
        self.lstm = nn.LSTM(embedding_size,
                            hidden_size,
                            num_layers=num_layers,
                            batch_first=True,
                            dropout=0 if num_layers == 1 else dropout)

    def forward(self, src_ids):
        """
        src_ids:[batch_size, seq_len]
        embedding:[batch_size, seq_len, embedding_size]
        out:[batch_size, seq_len, hidden_size]
        h_n:[num_layers * num_directions, batch_size, hidden_size]
        c_n:[num_layers * num_directions, batch_size, hidden_size]
        """
        embedding = self.embedding(src_ids)
        out, (h_n, c_n) = self.lstm(embedding)
        return out, (h_n, c_n)


# 解码器 Decoder
class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_size, hidden_size, num_layers=1, dropout=0.1):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=tgt_vocab.stoi[PAD])
        self.lstm = nn.LSTM(embedding_size,
                            hidden_size,
                            num_layers=num_layers,
                            batch_first=True,
                            dropout=0 if num_layers == 1 else dropout)
        self.fc = nn.Linear(hidden_size, vocab_size)

    # 据 previous token IDs 和当前 hidden state 生成下一个token的预测分布。
    def forward(self, y_prev_ids, hidden):
        """
        y_prev_ids:[batch_size, 1]
        hidden:
            h_n:[num_layers * num_directions, batch_size, hidden_size]
            c_n:[num_layers * num_directions, batch_size, hidden_size]
        embedding:[batch_size, seq_len, embedding_size]
        out:[batch_size, seq_len, hidden_size]
        logits:[batch_size, 1, vocab_size]
        """
        embedding = self.embedding(y_prev_ids)
        out, hidden = self.lstm(embedding, hidden)
        logits = self.fc(out)
        return logits, hidden


# 序列到序列 seq2seq
class seq2seq(nn.Module):
    def __init__(self, encoder, decoder, teacher_forcing_ratio=0.5):
        super(seq2seq, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.teacher_forcing_ratio = teacher_forcing_ratio

    def forward(self, src_ids, tgt_in_ids):
        B, T = tgt_in_ids.size()
        _, hidden = self.encoder(src_ids)

        # 从目标序列输入张量tgt_in_ids中提取第一个时间步的token，并将其维度扩展为(batch_size, 1)的形状。
        # tgt_in_ids[:, 0] - 取每个样本序列的第一个位置的token
        y_prev = tgt_in_ids[:, 0].unsqueeze(1)
        outputs = []

        # 平衡训练效果
        # teacher_forcing_ratio 控制真实值和预测值的使用比例
        # 既保证训练稳定性又模拟实际推理场景
        # 提高泛化能力
        # 让模型学会处理自己生成的"不完美"输入
        # 减少推理时的累积误差
        for t in range(1, T + 1):
            logits, hidden = self.decoder(y_prev, hidden)
            outputs.append(logits)
            # 根据 teacher_forcing_ratio 的概率决定是否使用“教师强制”策略（即使用真实标签作为下一时间步的输入）。
            # 如果启用“教师强制”，则使用目标输入中真实的下一个词作为下一时间步的输入。
            # 如果不使用“教师强制”，则使用模型自身预测的最可能词作为下一时间步的输入。
            use_teacher_forcing = random.random() < self.teacher_forcing_ratio
            # 最后一次循环(t = T)只计算输出，不更新输入
            if t < T:
                if use_teacher_forcing:
                    y_prev = tgt_in_ids[:, t].unsqueeze(1)
                else:
                    # 对 logits 沿最后一个维度取最大值索引（即预测的词 ID），作为下一时间步的输入
                    y_prev = logits.argmax(dim=-1)
        return torch.cat(outputs, dim=1)
